import warnings
from typing import Dict, List, Optional, Tuple, Union

from lagent.llms.base_llm import AsyncLLMMixin, BaseLLM


class APITemplateParser:
    """
    API模型的提示模板解析器。
    专门用于处理API类型模型（如OpenAI API）的提示模板。

    Args:
        meta_template: API模型的元模板配置
    """

    def __init__(self, meta_template: Optional[Dict] = None):
        self.meta_template = meta_template
        # 检查元模板格式
        if meta_template:
            assert isinstance(meta_template, list)
            # 存储角色名到配置的映射
            self.roles: Dict[str, dict] = dict()
            for item in meta_template:
                assert isinstance(item, dict)
                assert item['role'] not in self.roles, '角色名必须唯一!'
                self.roles[item['role']] = item.copy()

    def __call__(self, dialog: List[Union[str, List]]):
        """
        解析提示模板，并在需要时用元模板包装。
        对于设置了元模板且输入为列表的情况，返回完整的对话历史。
        每个对话项的格式为: {'role': 'user', 'content': '...'}

        Args:
            dialog: 中间提示模板（可能需要被元模板包装）

        Returns:
            处理后的提示或对话列表
        """
        assert isinstance(dialog, (str, list))
        if isinstance(dialog, str):
            return dialog

        if self.meta_template:
            prompt = list()
            # 控制是否继续生成提示
            generate = True
            for i, item in enumerate(dialog):
                if not generate:
                    break
                if isinstance(item, str):
                    if item.strip():
                        warnings.warn('API模型会忽略提示模板中的非空字符串。')
                else:
                    api_prompts = self._prompt2api(item)
                    prompt.append(api_prompts)

            # 合并连续的相同角色的提示
            new_prompt = list([prompt[0]])
            last_role = prompt[0]['role']
            for item in prompt[1:]:
                if item['role'] == last_role:
                    new_prompt[-1]['content'] += '\n' + item['content']
                else:
                    last_role = item['role']
                    new_prompt.append(item)
            prompt = new_prompt

        else:
            # 不使用元模板的简单处理
            prompt = ''
            last_sep = ''
            for item in dialog:
                if isinstance(item, str):
                    if item:
                        prompt += last_sep + item
                elif item.get('content', ''):
                    prompt += last_sep + item.get('content', '')
                last_sep = '\n'
        return prompt

    def _prompt2api(self, prompts: Union[List, str]) -> Tuple[str, bool]:
        """
        将提示转换为API格式。

        Args:
            prompts: 要转换的提示

        Returns:
            转换后的API格式提示
        """
        if isinstance(prompts, str):
            return prompts
        elif isinstance(prompts, dict):
            api_role = self._role2api_role(prompts)
            return api_role

        res = []
        for prompt in prompts:
            if isinstance(prompt, str):
                raise TypeError('API模型不允许混合使用没有明确角色的字符串!')
            else:
                api_role = self._role2api_role(prompt)
                res.append(api_role)
        return res

    def _role2api_role(self, role_prompt: Dict) -> Tuple[str, bool]:
        """
        将角色提示转换为API角色格式。

        Args:
            role_prompt: 角色提示配置

        Returns:
            API格式的角色配置
        """
        merged_prompt = self.roles[role_prompt['role']]
        if merged_prompt.get('fallback_role'):
            merged_prompt = self.roles[self.roles[
                merged_prompt['fallback_role']]]
        res = role_prompt.copy()
        res['role'] = merged_prompt['api_role']
        res['content'] = merged_prompt.get('begin', '')
        res['content'] += role_prompt.get('content', '')
        res['content'] += merged_prompt.get('end', '')
        return res


class BaseAPILLM(BaseLLM):
    """
    API模型包装器的基类。

    Args:
        model_type: 模型类型
        retry: API调用失败时的重试次数，默认2次
        meta_template: 模型的元提示模板（如果需要）
    """

    is_api: bool = True

    def __init__(self,
                 model_type: str,
                 retry: int = 2,
                 template_parser: 'APITemplateParser' = APITemplateParser,
                 meta_template: Optional[Dict] = None,
                 *,
                 max_new_tokens: int = 512,
                 top_p: float = 0.8,
                 top_k: int = 40,
                 temperature: float = 0.8,
                 repetition_penalty: float = 0.0,
                 stop_words: Union[List[str], str] = None):
        """初始化API模型参数"""
        self.model_type = model_type
        self.meta_template = meta_template
        self.retry = retry
        if template_parser:
            self.template_parser = template_parser(meta_template)

        # 处理停止词
        if isinstance(stop_words, str):
            stop_words = [stop_words]
            
        # 生成参数配置
        self.gen_params = dict(
            max_new_tokens=max_new_tokens,  # 最大生成长度
            top_p=top_p,  # 累积概率阈值
            top_k=top_k,  # 保留最高概率的k个token
            temperature=temperature,  # 采样温度
            repetition_penalty=repetition_penalty,  # 重复惩罚系数
            stop_words=stop_words,  # 停止词列表
            skip_special_tokens=False  # 是否跳过特殊token
        )


class AsyncBaseAPILLM(AsyncLLMMixin, BaseAPILLM):
    """异步API模型基类，结合异步接口和API基础功能"""
    pass
